package com.agentframework.core.memory;

import com.agentframework.core.action.Action;
import com.agentframework.core.agent.Agent;
import com.agentframework.core.plan.WritePlanAction;
import org.jetbrains.annotations.NotNull;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.memory.ChatMemoryRepository;
import org.springframework.ai.chat.memory.InMemoryChatMemoryRepository;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.util.Assert;

import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Set;

import static com.agentframework.core.agent.Agent.*;

/**
 * @author wangjun
 **/
public class AgentInMemoryChatMemory implements ChatMemory {

    protected static final Logger logger = LoggerFactory.getLogger(AgentInMemoryChatMemory.class);

    private Agent agent;

    private static final int DEFAULT_MAX_MESSAGES = 20;

    private final ChatMemoryRepository chatMemoryRepository;

    private final int maxMessages;

    private AgentInMemoryChatMemory(ChatMemoryRepository chatMemoryRepository, int maxMessages) {
        Assert.notNull(chatMemoryRepository, "chatMemoryRepository cannot be null");
        Assert.isTrue(maxMessages > 0, "maxMessages must be greater than 0");
        this.chatMemoryRepository = chatMemoryRepository;
        this.maxMessages = maxMessages;
    }

    @Override
    public void add(@NotNull String conversationId, @NotNull List<Message> messages) {
        Assert.hasText(conversationId, "conversationId cannot be null or empty");
        Assert.notNull(messages, "messages cannot be null");
        Assert.noNullElements(messages, "messages cannot contain null elements");

        if (agent != null) {
            // 模型排任务的消息会给对话造成干扰，不需要放到记忆中
            Action todo = agent.getContext().getTodo();
            if (todo instanceof WritePlanAction) {
                return;
            }
        }

        List<Message> memoryMessages = this.chatMemoryRepository.findByConversationId(conversationId);
        List<Message> processedMessages = process(memoryMessages, messages);
        this.chatMemoryRepository.saveAll(conversationId, processedMessages);
    }

    @NotNull
    @Override
    public List<Message> get(@NotNull String conversationId) {
        Assert.hasText(conversationId, "conversationId cannot be null or empty");
        List<Message> historyMessages = this.chatMemoryRepository.findByConversationId(conversationId);

//        if (!historyMessages.isEmpty()) {
//            logger.debug(">>>>>>>>>> agent {} all historyMessages in env print: \n >>> [ah] {}",
//                    this.agent,
//                    historyMessages.stream().map(Object::toString).collect(Collectors.joining("\n >>> [ah] ")));
//        }

        historyMessages = historyMessages.stream().filter(m -> {
            var msgFrom = m.getMetadata().get(MSG_FROM);
            var sendTo = m.getMetadata().get(SEND_TO);
            var viewTo = m.getMetadata().get(VIEW_TO);
            return (msgFrom != null && agent.getContext().getWatch().contains(msgFrom.toString()))
                    || (sendTo != null && ((Set<?>) sendTo).contains(this.agent.getAgentId()))
                    || (viewTo != null && ((Set<?>) viewTo).contains(this.agent.getAgentId()));
        }).toList();

//        if (!historyMessages.isEmpty()) {
//            logger.debug(">>>>>>>>>> agent {}, self historyMessage print: \n >>> [sh] {}",
//                    this.agent,
//                    historyMessages.stream().map(Object::toString).collect(Collectors.joining("\n >>> [sh] ")));
//        }

        return historyMessages;
    }

    @Override
    public void clear(@NotNull String conversationId) {
        Assert.hasText(conversationId, "conversationId cannot be null or empty");
        this.chatMemoryRepository.deleteByConversationId(conversationId);
    }

    private List<Message> process(List<Message> memoryMessages, List<Message> newMessages) {
        List<Message> processedMessages = new ArrayList<>();

        Set<Message> memoryMessagesSet = new HashSet<>(memoryMessages);
        boolean hasNewSystemMessage = newMessages.stream()
                .filter(SystemMessage.class::isInstance)
                .anyMatch(message -> !memoryMessagesSet.contains(message));

        memoryMessages.stream()
                .filter(message -> !(hasNewSystemMessage && message instanceof SystemMessage))
                .forEach(processedMessages::add);

        processedMessages.addAll(newMessages);

        if (processedMessages.size() <= this.maxMessages) {
            return processedMessages;
        }

        int messagesToRemove = processedMessages.size() - this.maxMessages;

        List<Message> trimmedMessages = new ArrayList<>();
        int removed = 0;
        for (Message message : processedMessages) {
            if (message instanceof SystemMessage || removed >= messagesToRemove) {
                trimmedMessages.add(message);
            }
            else {
                removed++;
            }
        }

        return trimmedMessages;
    }

    public static Builder builder() {
        return new Builder();
    }

    public static final class Builder {

        private ChatMemoryRepository chatMemoryRepository;

        private int maxMessages = DEFAULT_MAX_MESSAGES;

        private Builder() {
        }

        public Builder chatMemoryRepository(ChatMemoryRepository chatMemoryRepository) {
            this.chatMemoryRepository = chatMemoryRepository;
            return this;
        }

        public Builder maxMessages(int maxMessages) {
            this.maxMessages = maxMessages;
            return this;
        }

        public AgentInMemoryChatMemory build() {
            if (this.chatMemoryRepository == null) {
                this.chatMemoryRepository = new InMemoryChatMemoryRepository();
            }
            return new AgentInMemoryChatMemory(this.chatMemoryRepository, this.maxMessages);
        }

    }

    public AgentInMemoryChatMemory setAgent(Agent agent) {
        this.agent = agent;
        return this;
    }

}
